#include "chdb.h"
#include "chdb-internal.h"
#include "ChunkCollectorOutputFormat.h"
#include "LocalServer.h"

#include <cstddef>
#include <cstring>
#include <ChdbClient.h>
#include <EmbeddedServer.h>
#if USE_PYTHON
#    include <PythonTableCache.h>
#endif
#include <Common/MemoryTracker.h>
#include <Common/ThreadStatus.h>

#if defined(USE_MUSL) && defined(__aarch64__)
void chdb_musl_compile_stub(int arg)
{
    jmp_buf buf1;
    sigjmp_buf buf2;

    setjmp(buf1);
    sigsetjmp(buf2, arg);
}
#endif

#if USE_JEMALLOC
#    include <Common/memory.h>
#endif

#ifdef CHDB_STATIC_LIBRARY_BUILD
/// Force reference to ensure function registration object is linked
namespace DB
{
    extern void * ForceStaticRegistrationObjects();
}

[[maybe_unused]] void * force_link_function_references = DB::ForceStaticRegistrationObjects();
#endif

namespace CHDB
{

#if !USE_PYTHON
extern "C"
{
    extern chdb_state chdb_arrow_scan(chdb_connection, const char *, chdb_arrow_stream);
}

[[maybe_unused]] void * force_link_arrow_functions[] = {
    reinterpret_cast<void*>(chdb_arrow_scan)
};
#endif

// used only in pyEntryClickHouseLocal
static std::mutex CHDB_MUTEX;

static local_result_v2 * convert2LocalResultV2(QueryResult * query_result)
{
    auto * local_result = new local_result_v2();
    auto * materialized_query_result = static_cast<MaterializedQueryResult *>(query_result);

    if (!materialized_query_result)
    {
        String error = "Query processing failed";
        local_result->error_message = new char[error.size() + 1];
        std::memcpy(local_result->error_message, error.c_str(), error.size() + 1);
    }
    else if (!materialized_query_result->getError().empty())
    {
        const String & error = materialized_query_result->getError();
        local_result->error_message = new char[error.size() + 1];
        std::memcpy(local_result->error_message, error.c_str(), error.size() + 1);
    }
    else if (!materialized_query_result->result_buffer)
    {
        local_result->rows_read = materialized_query_result->rows_read;
        local_result->bytes_read = materialized_query_result->bytes_read;
        local_result->elapsed = materialized_query_result->elapsed;
    }
    else
    {
        local_result->len = materialized_query_result->result_buffer->size();
        local_result->buf = materialized_query_result->result_buffer->data();
        local_result->_vec = materialized_query_result->result_buffer.release();
        local_result->rows_read = materialized_query_result->rows_read;
        local_result->bytes_read = materialized_query_result->bytes_read;
        local_result->elapsed = materialized_query_result->elapsed;
    }

    return local_result;
}

static local_result_v2 * createErrorLocalResultV2(const String & error)
{
    auto * local_result = new local_result_v2();
    local_result->error_message = new char[error.size() + 1];
    std::memcpy(local_result->error_message, error.c_str(), error.size() + 1);
    return local_result;
}

std::unique_ptr<MaterializedQueryResult> pyEntryClickHouseLocal(int argc, char ** argv)
{
    try
    {
        std::lock_guard<std::mutex> lock(CHDB_MUTEX);
        DB::ThreadStatus thread_status;
        DB::LocalServer app;
        app.init(argc, argv);
        int ret = app.run();
        if (ret == 0)
        {
            return std::make_unique<MaterializedQueryResult>(
                ResultBuffer(app.stealQueryOutputVector()), app.getElapsedTime(), app.getProcessedRows(), app.getProcessedBytes(), 0, 0);
        }
        else
        {
            return std::make_unique<MaterializedQueryResult>(app.getErrorMsg());
        }
    }
    catch (const DB::Exception & e)
    {
        // wrap the error message into a new std::exception
        throw std::domain_error(DB::getExceptionMessage(e, false));
    }
    catch (const boost::program_options::error & e)
    {
        throw std::invalid_argument("Bad arguments: " + std::string(e.what()));
    }
    catch (...)
    {
        throw std::domain_error(DB::getCurrentExceptionMessage(true));
    }
}

const static std::string empty_string;
const std::string & chdb_result_error_string(chdb_result * result)
{
    if (!result)
        return empty_string;

    auto * query_result = reinterpret_cast<QueryResult *>(result);
    return query_result->getError();
}

const std::string & chdb_streaming_result_error_string(chdb_streaming_result * result)
{
    if (!result)
        return empty_string;

    auto * stream_query_result = reinterpret_cast<StreamQueryResult *>(result);
    return stream_query_result->getError();
}

chdb_connection * connect_chdb_with_exception(int argc, char ** argv)
{
    try
    {
        DB::ThreadStatus thread_status;
        auto server = DB::EmbeddedServer::getInstance(argc, argv);
        auto client = DB::ChdbClient::create(server);
        if (!client)
        {
            throw DB::Exception(DB::ErrorCodes::BAD_ARGUMENTS, "Failed to create ChdbClient");
        }

        auto * conn = new chdb_conn();
        conn->server = client.release();
        conn->connected = true;
        auto ** conn_ptr = new chdb_conn *(conn);
        return reinterpret_cast<chdb_connection *>(conn_ptr);
    }
    catch (const DB::Exception & e)
    {
        throw DB::Exception(DB::ErrorCodes::BAD_ARGUMENTS, "Failed to create connection: {}", DB::getExceptionMessage(e, false));
    }
    catch (const std::exception & e)
    {
        throw std::domain_error(std::string("Connection failed: ") + e.what());
    }
    catch (...)
    {
        throw std::domain_error(DB::getCurrentExceptionMessage(true));
    }
}

#if USE_PYTHON
void cachePythonTablesFromQuery(chdb_conn * conn, const std::string & query_str)
{
    if (!conn || !conn->server || !conn->connected)
        return;
    auto * client = reinterpret_cast<DB::ChdbClient *>(conn->server);
    client->findQueryableObjFromPyCache(query_str);
}
#endif

} // namespace CHDB

using namespace CHDB;

local_result * query_stable(int argc, char ** argv)
{
    auto query_result = pyEntryClickHouseLocal(argc, argv);
    if (!query_result->getError().empty() || query_result->result_buffer == nullptr)
        return nullptr;

    local_result * res = new local_result;
    res->len = query_result->result_buffer->size();
    res->buf = query_result->result_buffer->data();
    res->_vec = query_result->result_buffer.release();
    res->rows_read = query_result->rows_read;
    res->bytes_read = query_result->bytes_read;
    res->elapsed = query_result->elapsed;
    return res;
}

void free_result(local_result * result)
{
    if (!result)
    {
        return;
    }
    if (result->_vec)
    {
        std::vector<char> * vec = reinterpret_cast<std::vector<char> *>(result->_vec);
        delete vec;
        result->_vec = nullptr;
    }
    delete result;
}

local_result_v2 * query_stable_v2(int argc, char ** argv)
{
    // pyEntryClickHouseLocal may throw some serious exceptions, although it's not likely
    // to happen in the context of clickhouse-local. we catch them here and return an error
    local_result_v2 * res = nullptr;
    try
    {
        auto query_result = pyEntryClickHouseLocal(argc, argv);
        return convert2LocalResultV2(query_result.get());
    }
    catch (const std::exception & e)
    {
        res = new local_result_v2();
        res->error_message = new char[strlen(e.what()) + 1];
        std::strcpy(res->error_message, e.what());
    }
    catch (...)
    {
        res = new local_result_v2();
        const char * unknown_exception_msg = "Unknown exception";
        size_t len = std::strlen(unknown_exception_msg) + 1;
        res->error_message = new char[len];
        std::strcpy(res->error_message, unknown_exception_msg);
    }

    return res;
}

void free_result_v2(local_result_v2 * result)
{
    if (!result)
        return;

    delete reinterpret_cast<std::vector<char> *>(result->_vec);
    delete[] result->error_message;
    delete result;
}

chdb_conn ** connect_chdb(int argc, char ** argv)
{
    auto * connection = chdb_connect(argc, argv);
    if (!connection)
    {
        return nullptr;
    }
    return reinterpret_cast<chdb_conn **>(connection);
}

void close_conn(chdb_conn ** conn)
{
    if (!conn || !*conn)
        return;

    try
    {
        if ((*conn)->connected && (*conn)->server)
        {
            auto * client = static_cast<DB::ChdbClient *>((*conn)->server);
            delete client;
            (*conn)->server = nullptr;
        }

        (*conn)->connected = false;
        delete *conn;
        *conn = nullptr;
    }
    catch (...)
    {
        DB::tryLogCurrentException(__PRETTY_FUNCTION__);
    }
}

struct local_result_v2 * query_conn(chdb_conn * conn, const char * query, const char * format)
{
    return query_conn_n(conn, query, query ? std::strlen(query) : 0, format, format ? std::strlen(format) : 0);
}

struct local_result_v2 * query_conn_n(struct chdb_conn * conn, const char * query, size_t query_len, const char * format, size_t format_len)
{
    if (!checkConnectionValidity(conn))
        return createErrorLocalResultV2("Invalid or closed connection");

    try
    {
        auto * client = static_cast<DB::ChdbClient *>(conn->server);
        auto query_result = client->executeMaterializedQuery(query, query_len, format, format_len);
        return convert2LocalResultV2(query_result.get());
    }
    catch (const std::exception & e)
    {
        return createErrorLocalResultV2(std::string("Error: ") + e.what());
    }
    catch (...)
    {
        return createErrorLocalResultV2(DB::getCurrentExceptionMessage(true));
    }
}

chdb_streaming_result * query_conn_streaming(chdb_conn * conn, const char * query, const char * format)
{
    return query_conn_streaming_n(conn, query, query ? std::strlen(query) : 0, format, format ? std::strlen(format) : 0);
}

chdb_streaming_result *
query_conn_streaming_n(struct chdb_conn * conn, const char * query, size_t query_len, const char * format, size_t format_len)
{
    if (!checkConnectionValidity(conn))
    {
        auto * result = new StreamQueryResult("Invalid or closed connection");
        return reinterpret_cast<chdb_streaming_result *>(result);
    }

    try
    {
        // Use the client from the connection
        auto * client = static_cast<DB::ChdbClient *>(conn->server);
        auto query_result = client->executeStreamingInit(query, query_len, format, format_len);

        if (!query_result)
        {
            auto * result = new StreamQueryResult("Query processing failed");
            return reinterpret_cast<chdb_streaming_result *>(result);
        }

        return reinterpret_cast<chdb_streaming_result *>(query_result.release());
    }
    catch (const std::exception & e)
    {
        auto * result = new StreamQueryResult(std::string("Error: ") + e.what());
        return reinterpret_cast<chdb_streaming_result *>(result);
    }
    catch (...)
    {
        auto * result = new StreamQueryResult(DB::getCurrentExceptionMessage(true));
        return reinterpret_cast<chdb_streaming_result *>(result);
    }
}

const char * chdb_streaming_result_error(chdb_streaming_result * result)
{
    if (!result)
        return nullptr;

    auto * stream_query_result = reinterpret_cast<StreamQueryResult *>(result);

    const auto & error_message = stream_query_result->getError();
    if (!error_message.empty())
        return error_message.c_str();

    return nullptr;
}

local_result_v2 * chdb_streaming_fetch_result(chdb_conn * conn, chdb_streaming_result * result)
{
    if (!checkConnectionValidity(conn))
        return createErrorLocalResultV2("Invalid or closed connection");

    if (!result)
        return createErrorLocalResultV2("Invalid streaming result");

    try
    {
        auto * client = static_cast<DB::ChdbClient *>(conn->server);
        if (!client->hasStreamingQuery())
            return createErrorLocalResultV2("No active streaming query");
        auto query_result = client->executeStreamingIterate(result, false);
        if (!query_result)
            return createErrorLocalResultV2("Failed to fetch streaming results");
        auto * local_result = convert2LocalResultV2(query_result.release());
        return local_result;
    }
    catch (const std::exception & e)
    {
        return createErrorLocalResultV2(std::string("Error fetching streaming results: ") + e.what());
    }
    catch (...)
    {
        return createErrorLocalResultV2(std::string("Unknown error fetching streaming results: ") + DB::getCurrentExceptionMessage(true));
    }
}

void chdb_streaming_cancel_query(chdb_conn * conn, chdb_streaming_result * result)
{
    if (!checkConnectionValidity(conn))
        return;

    if (!result)
        return;

    try
    {
        auto * client = static_cast<DB::ChdbClient *>(conn->server);
        client->cancelStreamingQuery(result);
    }
    catch (...)
    {
        DB::tryLogCurrentException(__PRETTY_FUNCTION__);
    }
    /// Note: The result object should be freed by chdb_destroy_result(), not here
}

void chdb_destroy_result(chdb_streaming_result * result)
{
    if (!result)
        return;

    auto * stream_query_result = reinterpret_cast<StreamQueryResult *>(result);

    delete stream_query_result;
}

/// ============== New API Implementation ==============

chdb_connection * chdb_connect(int argc, char ** argv)
{
    try
    {
        return connect_chdb_with_exception(argc, argv);
    }
    catch (const DB::Exception & e)
    {
        LOG_ERROR(&Poco::Logger::get("EmbeddedServer"), "Connection failed with DB::Exception: {}", DB::getExceptionMessage(e, false));
        return nullptr;
    }
    catch (const boost::program_options::error & e)
    {
        LOG_ERROR(&Poco::Logger::get("EmbeddedServer"), "Connection failed with bad arguments: {}", e.what());
        return nullptr;
    }
    catch (const Poco::Exception & e)
    {
        LOG_ERROR(&Poco::Logger::get("EmbeddedServer"), "Connection failed with Poco::Exception: {}", e.displayText());
        return nullptr;
    }
    catch (...)
    {
        LOG_ERROR(
            &Poco::Logger::get("EmbeddedServer"), "Connection failed with unknown exception: {}", DB::getCurrentExceptionMessage(true));
        return nullptr;
    }
}

void chdb_close_conn(chdb_connection * conn)
{
    if (!conn || !*conn)
        return;

    auto * connection = reinterpret_cast<chdb_conn **>(conn);

    close_conn(connection);
}

chdb_result * chdb_query(chdb_connection conn, const char * query, const char * format)
{
    return chdb_query_n(conn, query, query ? std::strlen(query) : 0, format, format ? std::strlen(format) : 0);
}

chdb_result * chdb_query_n(chdb_connection conn, const char * query, size_t query_len, const char * format, size_t format_len)
{
    if (!conn)
    {
        auto * result = new MaterializedQueryResult("Unexpected null connection");
        return reinterpret_cast<chdb_result *>(result);
    }

    auto * connection = reinterpret_cast<chdb_conn *>(conn);
    if (!checkConnectionValidity(connection))
    {
        auto * result = new MaterializedQueryResult("Invalid or closed connection");
        return reinterpret_cast<chdb_result *>(result);
    }

    try
    {
        // Use the client from the connection
        auto * client = static_cast<DB::ChdbClient *>(connection->server);
        auto query_result = client->executeMaterializedQuery(query, query_len, format, format_len);

        return reinterpret_cast<chdb_result *>(query_result.release());
    }
    catch (const std::exception & e)
    {
        auto * result = new MaterializedQueryResult(std::string("Error: ") + e.what());
        return reinterpret_cast<chdb_result *>(result);
    }
    catch (...)
    {
        auto * result = new MaterializedQueryResult(DB::getCurrentExceptionMessage(true));
        return reinterpret_cast<chdb_result *>(result);
    }
}

chdb_result * chdb_query_cmdline(int argc, char ** argv)
{
    MaterializedQueryResult * result = nullptr;
    try
    {
        auto query_result = pyEntryClickHouseLocal(argc, argv);

        return reinterpret_cast<chdb_result *>(query_result.release());
    }
    catch (const std::exception & e)
    {
        result = new MaterializedQueryResult(e.what());
    }
    catch (...)
    {
        result = new MaterializedQueryResult("Unknown exception");
    }

    return reinterpret_cast<chdb_result *>(result);
}

chdb_result * chdb_stream_query(chdb_connection conn, const char * query, const char * format)
{
    return chdb_stream_query_n(conn, query, query ? std::strlen(query) : 0, format, format ? std::strlen(format) : 0);
}

chdb_result * chdb_stream_query_n(chdb_connection conn, const char * query, size_t query_len, const char * format, size_t format_len)
{
    if (!conn)
    {
        auto * result = new StreamQueryResult("Unexpected null connection");
        return reinterpret_cast<chdb_result *>(result);
    }

    auto * connection = reinterpret_cast<chdb_conn *>(conn);
    if (!checkConnectionValidity(connection))
    {
        auto * result = new StreamQueryResult("Invalid or closed connection");
        return reinterpret_cast<chdb_result *>(result);
    }

    try
    {
        auto * client = static_cast<DB::ChdbClient *>(connection->server);
        auto query_result = client->executeStreamingInit(query, query_len, format, format_len);

        if (!query_result)
        {
            auto * result = new StreamQueryResult("Query processing failed");
            return reinterpret_cast<chdb_result *>(result);
        }

        return reinterpret_cast<chdb_result *>(query_result.release());
    }
    catch (const std::exception & e)
    {
        auto * result = new StreamQueryResult(std::string("Error: ") + e.what());
        return reinterpret_cast<chdb_result *>(result);
    }
    catch (...)
    {
        auto * result = new StreamQueryResult(DB::getCurrentExceptionMessage(true));
        return reinterpret_cast<chdb_result *>(result);
    }
}

chdb_result * chdb_stream_fetch_result(chdb_connection conn, chdb_result * result)
{
    if (!conn)
    {
        auto * query_result = new MaterializedQueryResult("Unexpected null connection");
        return reinterpret_cast<chdb_result *>(query_result);
    }

    if (!result)
    {
        auto * query_result = new MaterializedQueryResult("Unexpected null result");
        return reinterpret_cast<chdb_result *>(query_result);
    }

    auto * connection = reinterpret_cast<chdb_conn *>(conn);
    if (!checkConnectionValidity(connection))
    {
        auto * query_result = new MaterializedQueryResult("Invalid or closed connection");
        return reinterpret_cast<chdb_result *>(query_result);
    }

    try
    {
        auto * client = static_cast<DB::ChdbClient *>(connection->server);
        if (!client->hasStreamingQuery())
            return reinterpret_cast<chdb_result *>(new MaterializedQueryResult("No active streaming query"));
        auto * stream_result = reinterpret_cast<StreamQueryResult *>(result);
        auto query_result = client->executeStreamingIterate(stream_result, false);
        if (!query_result)
            return reinterpret_cast<chdb_result *>(new MaterializedQueryResult("Failed to fetch streaming results"));

        return reinterpret_cast<chdb_result *>(query_result.release());
    }
    catch (const std::exception & e)
    {
        auto * query_result = new MaterializedQueryResult(std::string("Error: ") + e.what());
        return reinterpret_cast<chdb_result *>(query_result);
    }
    catch (...)
    {
        auto * query_result = new MaterializedQueryResult(DB::getCurrentExceptionMessage(true));
        return reinterpret_cast<chdb_result *>(query_result);
    }
}

void chdb_stream_cancel_query(chdb_connection conn, chdb_result * result)
{
    if (!result || !conn)
        return;

    auto * connection = reinterpret_cast<chdb_conn *>(conn);
    if (!checkConnectionValidity(connection))
        return;

    try
    {
        auto * client = static_cast<DB::ChdbClient *>(connection->server);
        auto * stream_result = reinterpret_cast<StreamQueryResult *>(result);
        client->cancelStreamingQuery(stream_result);
    }
    catch (...)
    {
        DB::tryLogCurrentException(__PRETTY_FUNCTION__);
    }
    /// Note: The result object should be freed by chdb_destroy_query_result(), not here
}

void chdb_destroy_query_result(chdb_result * result)
{
    if (!result)
        return;

    auto * query_result = reinterpret_cast<QueryResult *>(result);
    delete query_result;
}

char * chdb_result_buffer(chdb_result * result)
{
    if (!result)
        return nullptr;

    auto * query_result = reinterpret_cast<QueryResult *>(result);

    if (query_result->getType() == QueryResultType::RESULT_TYPE_MATERIALIZED)
    {
        auto * materialized_result = reinterpret_cast<MaterializedQueryResult *>(result);
        return materialized_result->result_buffer ? materialized_result->result_buffer->data() : nullptr;
    }

    return nullptr;
}

size_t chdb_result_length(chdb_result * result)
{
    if (!result)
        return 0;

    auto * query_result = reinterpret_cast<QueryResult *>(result);
    if (query_result->getType() == QueryResultType::RESULT_TYPE_MATERIALIZED)
    {
        auto * materialized_result = reinterpret_cast<MaterializedQueryResult *>(result);
        return materialized_result->result_buffer ? materialized_result->result_buffer->size() : 0;
    }

    return 0;
}

double chdb_result_elapsed(chdb_result * result)
{
    if (!result)
        return 0.0;

    auto * query_result = reinterpret_cast<QueryResult *>(result);

    if (query_result->getType() == QueryResultType::RESULT_TYPE_MATERIALIZED)
    {
        auto * materialized_result = reinterpret_cast<MaterializedQueryResult *>(result);
        return materialized_result->elapsed;
    }
    return 0.0;
}

uint64_t chdb_result_rows_read(chdb_result * result)
{
    if (!result)
        return 0;

    auto * query_result = reinterpret_cast<QueryResult *>(result);

    if (query_result->getType() == QueryResultType::RESULT_TYPE_MATERIALIZED)
    {
        auto * materialized_result = reinterpret_cast<MaterializedQueryResult *>(result);
        return materialized_result->rows_read;
    }

    return 0;
}

uint64_t chdb_result_bytes_read(chdb_result * result)
{
    if (!result)
        return 0;

    auto * query_result = reinterpret_cast<QueryResult *>(result);

    if (query_result->getType() == QueryResultType::RESULT_TYPE_MATERIALIZED)
    {
        auto * materialized_result = reinterpret_cast<MaterializedQueryResult *>(result);
        return materialized_result->bytes_read;
    }

    return 0;
}

uint64_t chdb_result_storage_rows_read(chdb_result * result)
{
    if (!result)
        return 0;

    auto * query_result = reinterpret_cast<QueryResult *>(result);

    if (query_result->getType() == QueryResultType::RESULT_TYPE_MATERIALIZED)
    {
        auto * materialized_result = reinterpret_cast<MaterializedQueryResult *>(result);
        return materialized_result->storage_rows_read;
    }

    return 0;
}

uint64_t chdb_result_storage_bytes_read(chdb_result * result)
{
    if (!result)
        return 0;

    auto * query_result = reinterpret_cast<QueryResult *>(result);

    if (query_result->getType() == QueryResultType::RESULT_TYPE_MATERIALIZED)
    {
        auto * materialized_result = reinterpret_cast<MaterializedQueryResult *>(result);
        return materialized_result->storage_bytes_read;
    }

    return 0;
}

const char * chdb_result_error(chdb_result * result)
{
    if (!result)
        return nullptr;

    auto * query_result = reinterpret_cast<QueryResult *>(result);

    if (query_result->getError().empty())
        return nullptr;

    return query_result->getError().c_str();
}
